# -*- coding: utf-8 -*-
"""
🚀 高性能异步 TCP 压测脚本（直方图统计 | 无内存增长 | 长期运行安全）
Author: chen xi
Date: 2025/9/24
功能：
- 异步并发发送 TCP 请求
- 每 10 秒输出实时 QPS、延迟、成功率
- 最终输出汇总报告（基于直方图计算 P95/P99）
- 使用直方图代替原始延迟存储，避免内存泄漏
"""

import asyncio
import time
import sys
from typing import Tuple, DefaultDict
from collections import defaultdict

# =================== 配置项 ===================
SERVER_IP = '47.107.56.10'
SERVER_PORT = 50002
MESSAGE = b'0123456789'           # 发送的数据
NUM_WORKERS = 500                  # 并发协程数（建议调高）
DURATION = 10 * 60 * 60 * 24                # 测试总时长（秒），示例：10分钟
STAT_INTERVAL = 10                # 统计间隔（秒）
TIMEOUT = 5                       # 网络超时（秒）
RECV_SIZE = 1024                  # 接收缓冲区大小
# =============================================

# 直方图配置：1ms 一桶，最大支持 5000ms（5秒）
BUCKET_SIZE: float = 1.0          # 每个桶的宽度（ms）
MAX_BUCKET: int = 5000             # 最大桶编号（对应 5000ms）

# 全局统计（整数型，无内存增长）
total_success = 0
total_fail = 0
total_bytes_received = 0

# 延迟直方图：key=桶编号, value=计数
latency_histogram: DefaultDict[int, int] = defaultdict(int)

# 当前窗口统计
current_success = 0
current_fail = 0
current_latencies = []  # 仅用于窗口内 P95 计算，会清空
current_bytes = 0

# 锁
lock = asyncio.Lock()

# 控制变量
running = True


async def tcp_request() -> Tuple[bool, float, int]:
    """
    执行一次 TCP 请求
    Returns: (success, latency_ms, bytes_received)
    """
    start_time = time.time()
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(SERVER_IP, SERVER_PORT),
            timeout=TIMEOUT
        )

        if MESSAGE:
            writer.write(MESSAGE)
            await writer.drain()

        response = await asyncio.wait_for(reader.read(RECV_SIZE), timeout=TIMEOUT)

        writer.close()
        await writer.wait_closed()

        latency_ms = (time.time() - start_time) * 1000
        return True, latency_ms, len(response)

    except asyncio.TimeoutError:
        return False, 0, 0
    except (ConnectionRefusedError, ConnectionResetError):
        return False, 0, 0
    except Exception as e:
        print(f"❌ 未知错误 [{type(e).__name__}]: {e}", file=sys.stderr)
        return False, 0, 0


def record_latency(latency_ms: float):
    """将延迟记录到直方图"""
    bucket = min(int(latency_ms / BUCKET_SIZE), MAX_BUCKET)
    latency_histogram[bucket] += 1


async def worker():
    global running, current_success, current_fail, current_latencies, current_bytes
    while running:
        success, latency_ms, bytes_recv = await tcp_request()

        async with lock:
            if success:
                current_success += 1
                current_latencies.append(latency_ms)
                current_bytes += bytes_recv
                # ✅ 记录到全局直方图
                record_latency(latency_ms)
            else:
                current_fail += 1

        await asyncio.sleep(0.01)  # 控制请求频率，避免过于密集


async def print_stats():
    """每 STAT_INTERVAL 秒打印一次统计"""
    global current_success, current_fail, current_latencies, current_bytes
    global total_success, total_fail, total_bytes_received

    start_time = time.time()
    last_print = start_time

    print(f"🚀 压测开始，目标: {SERVER_IP}:{SERVER_PORT} | 并发: {NUM_WORKERS} 协程 | 每 {STAT_INTERVAL} 秒输出一次")
    print("-" * 80)
    print(f"{'时间':<8} {'QPS':<8} {'平均延迟(ms)':<12} {'P95延迟(ms)':<12} {'成功率':<8} {'接收数据'}")
    print("-" * 80)

    while time.time() - start_time < DURATION:
        await asyncio.sleep(STAT_INTERVAL - (time.time() - last_print) % STAT_INTERVAL)

        async with lock:
            duration = time.time() - last_print
            qps = current_success / duration if duration > 0 else 0
            avg_lat = sum(current_latencies) / len(current_latencies) if current_latencies else 0
            p95_lat = 0
            if len(current_latencies) >= 20:
                try:
                    p95_lat = sorted(current_latencies)[int(0.95 * len(current_latencies))]
                except:
                    p95_lat = 0
            success_rate = f"{current_success / (current_success + current_fail) * 100:.1f}%" if (current_success + current_fail) > 0 else "100%"

            print(f"{int(time.time() - start_time):>8}s "
                  f"{qps:>8.1f} "
                  f"{avg_lat:>12.2f} "
                  f"{p95_lat:>12.2f} "
                  f"{success_rate:>8} "
                  f"{current_bytes:,} B")

            # ✅ 累加到总统计（仅计数，无数据存储）
            total_success += current_success
            total_fail += current_fail
            total_bytes_received += current_bytes

            # 重置当前窗口
            current_success = 0
            current_fail = 0
            current_latencies.clear()
            current_bytes = 0

        last_print = time.time()


def calculate_percentile(hist: DefaultDict[int, int], percentile: float) -> float:
    """根据直方图计算指定分位数（如 P95）"""
    total = sum(hist.values())
    if total == 0:
        return 0.0
    threshold = total * percentile / 100
    count = 0
    for bucket in sorted(hist.keys()):
        count += hist[bucket]
        if count >= threshold:
            # 返回该桶的起始值（可优化为中点：(bucket + 0.5) * BUCKET_SIZE）
            return bucket * BUCKET_SIZE
    return MAX_BUCKET * BUCKET_SIZE


async def main():
    global running
    print("🚀 正在启动压测...", file=sys.stderr)

    # 启动统计任务
    stats_task = asyncio.create_task(print_stats())

    # 启动 workers
    workers = [asyncio.create_task(worker()) for _ in range(NUM_WORKERS)]
    print(f"✅ 已启动 {len(workers)} 个 worker 协程", file=sys.stderr)

    # 运行 DURATION 秒
    await asyncio.sleep(DURATION)

    # 停止
    running = False
    print("🛑 正在停止所有 worker...", file=sys.stderr)
    await asyncio.gather(*workers, return_exceptions=True)
    await stats_task

    # =================== 最终报告 ===================
    total_time = DURATION
    total_requests = total_success + total_fail
    final_qps = total_success / total_time if total_time > 0 else 0
    throughput_mbps = (total_bytes_received * 8) / (total_time * 1024 * 1024) if total_time > 0 else 0

    # 从直方图计算全局延迟
    final_avg_lat = sum(latency * count for latency, count in ((k * BUCKET_SIZE, v) for k, v in latency_histogram.items())) / total_success if total_success > 0 else 0
    final_p95_lat = calculate_percentile(latency_histogram, 95)
    final_p99_lat = calculate_percentile(latency_histogram, 99)

    print("=" * 80)
    print("📊 压测完成！最终性能报告（基于直方图统计）")
    print("=" * 80)
    print(f"⏱  测试时长:        {total_time:.2f} 秒")
    print(f"🎯 总请求数:        {total_requests}")
    print(f"✅ 成功请求数:      {total_success}")
    print(f"❌ 失败请求数:      {total_fail}")
    print(f"🔥 平均 QPS:        {final_qps:.2f}")
    print(f"⚡ 平均响应时间:    {final_avg_lat:.2f} ms")
    print(f"📌 95% 响应时间:    {final_p95_lat:.2f} ms")
    print(f"📌 99% 响应时间:    {final_p99_lat:.2f} ms")
    print(f"📦 吞吐量:         {throughput_mbps:.2f} Mbps")
    print(f"📤 总接收字节数:    {total_bytes_received:,} 字节")
    print(f"📊 直方图桶数:      {len(latency_histogram)} (内存恒定)")
    print("=" * 80)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n👋 压测被用户中断", file=sys.stderr)
    except Exception as e:
        print(f"🚨 主程序异常: {e}", file=sys.stderr)